/***************************************
 * detect_rune(): algorithm retrieved from \
 *                  https://github.com/EinstainBohr/Windmills author: EinstainBohr
 *              
 * 
 * */

/*
TODO: find the proper light exposure
TODO: delta_angle in armor_detect(), to be modified
TODO: find a proper transition matrix and measurematrix, control matrix in order to measure the speed of the 
TODO: ROI method
*/
#define DEBUG_MODE
#include <armor_recog.hpp>

#include <algorithm>

using namespace std;
using namespace cv;

float formalize_angle(float angle)
{
    if (angle <= 0)
    {
        if (angle >= -45)
            return angle + 90;
        else if (angle < -45) //-90?
            return angle + 180;
    }
    else if (angle > 0)
    {
        if (angle > 45)
            return angle;
        else
            return angle + 90;
    }
    return angle;
}

Armor_recog::Armor_recog(int enemyColor, int sentry_mode)
{
    Armor_recog::enemyColor = enemyColor;
    Armor_recog::sentry_mode = sentry_mode;
    to_aim = Armor(); //initialize as empty Armor
    KF_initialize();
    if (enemyColor == RED)
    {
        threshold_color = THRESHOLD_MAX_COLOR_RED;
        threshold_brightness = BRIGHTNESS_THRES_RED;
    }
    else if (enemyColor == BLUE)
    {
        threshold_color = THRESHOLD_MAX_COLOR_BLUE;
        threshold_brightness = BRIGHTNESS_THRES_BLUE;
    }
}

bool Armor_recog::read_frame(cv::Mat frame)
{
    frame.copyTo(current_frame);
    now = std::chrono::steady_clock::now();
    if (current_frame.empty())
    {
        return false;
    }
    return true;
}

void Armor_recog::pre_process()
{
    Mat dilateElement = getStructuringElement(cv::MORPH_ELLIPSE, Size(3, 3));
    Mat thres_whole;
    Mat color;

    vector<Mat> splited;
    split(current_frame, splited);
    cvtColor(current_frame, thres_whole, cv::COLOR_BGR2GRAY);
    threshold(thres_whole, thres_whole, threshold_brightness, 255, cv::THRESH_BINARY); //binary-lize the image
    dilate(thres_whole, thres_whole, dilateElement);

#ifdef DEBUG_MODE
    imshow("origin", current_frame);
    imshow("thresh_whole", thres_whole);
#endif

    if (enemyColor == BLUE)
    {
        subtract(splited[0], splited[2], color); //subtract the Blue channel with Red channel
    }
    else
    {
        subtract(splited[2], splited[0], color);
#ifdef DEBUG_MODE
        imshow("BLUE", splited[0]);
        imshow("RED", splited[2]);
#endif
    }
    // imshow("subtracted1", color);
    threshold(color, color, threshold_color, 255, cv::THRESH_BINARY); //binary-lize to obtain the blue area
    // imshow("subtracted2", color);
    // dilate(color, color, dilateElement); //dilate, enlarge the bright area of the input image, which is the blue area
    // imshow("dilated", color);

    processed_frame = color & thres_whole; //use & logic operator to get rid of area that is bright but not the target color
    // dilate(processed_frame, processed_frame, dilateElement2);
#ifdef DEBUG_MODE
    imshow("color", color);
    imshow("result", processed_frame);
#endif
}

void Armor_recog::edge_detect()
{
    vector<vector<Point2i>> contours_max;
    vector<Vec4i> hierarchy;
    /* linear regression and classification, fit multiple points into one rotate Rectangle
    that is, find the boundaries*/
    findContours(processed_frame, contours_max, hierarchy, RETR_EXTERNAL, CHAIN_APPROX_SIMPLE);

    for (size_t i = 0; i < contours_max.size(); i++)
    {
        RotatedRect rRect;
        // if (contours_max[i].size() <= 5)
        // {
        rRect = minAreaRect(contours_max[i]);
        // }else
        // {
        //     rRect = fitEllipse(contours_max[i]);
        // }
        double rRect_length = max(rRect.size.width, rRect.size.height);
        double rRect_width = min(rRect.size.width, rRect.size.height);
        //fiter out those are not likely a RGB LED bar
        bool isWidthValid = (rRect_width > 1);
        bool isAProperRect = (rRect_length / rRect_width >= 1.1) && ((rRect_length / rRect_width) < 20);
        bool NotHorizontal1 = ((abs(rRect.angle) < 30) && (rRect.size.height > rRect.size.width));
        bool NotHorizontal2 = ((abs(rRect.angle) > 60) && (rRect.size.height < rRect.size.width));
        if (isWidthValid && (NotHorizontal1 || NotHorizontal2) && isAProperRect)
            edges_detected.push_back(rRect);
    }
}

bool Armor_recog::armor_detect()
{
    //     armor_detected.clear();
    if (edges_detected.size() < 2)
    {
        cout << "no target detected\n";
        return false;
    }
    else
    {
        for (int i = 0; i < edges_detected.size() - 1; i++)
        {
            //RotatedRect.points, the array start with the point with smallest y value, and goes clockwise
            // ,,, for most rotatedrects got from minAreaRect(bars),
            // it's the array start with the point with biggest y value(and smallest x value when y is the same), and goes clockwise
            Point2f pointsI[4];
            Point2f pointsJ[4];
            edges_detected[i].points(pointsI);
            double lengthI = max(edges_detected[i].size.width, edges_detected[i].size.height);
            double widthI = min(edges_detected[i].size.width, edges_detected[i].size.width);
            double angleI = formalize_angle(edges_detected[i].angle);
            // double angleI = abs(edges_detected[i].angle);
            double deltaAngle;
            double yminI = pointsI[0].y;
            double ymaxI = pointsI[0].y;
            for (size_t i = 1; i < 4; i++)
            {
                if (pointsI[i].y < yminI)
                    yminI = pointsI[i].y;
                if (pointsI[i].y > ymaxI)
                    ymaxI = pointsI[i].y;
            }

            for (int j = i + 1; j < edges_detected.size(); j++)
            {
                edges_detected[j].points(pointsJ);
                double lengthJ = max(edges_detected[j].size.width, edges_detected[j].size.height);
                double widthJ = min(edges_detected[j].size.width, edges_detected[j].size.width);
                double angleJ = formalize_angle(edges_detected[j].angle);
                // double angleJ = abs(edges_detected[j].angle);
                double distance = sqrtf(square(edges_detected[i].center.x - edges_detected[j].center.x) + square(edges_detected[i].center.y - edges_detected[j].center.y));
                double length_panel = (lengthI + lengthJ) / 2;
                double length_width_ratio_panel = distance / length_panel;
                double yminJ = pointsJ[0].y;
                double ymaxJ = pointsJ[0].y;

                for (size_t i = 1; i < 4; i++)
                {
                    if (pointsJ[i].y < yminJ)
                        yminJ = pointsJ[i].y;
                    if (pointsJ[i].y > ymaxJ)
                        ymaxJ = pointsJ[i].y;
                }

                deltaAngle = abs(angleI - angleJ);
                bool isTwoEdgeAtSameLevel = yminI < ymaxJ && yminJ < ymaxI;
                bool isDeltaAngleValid = deltaAngle < DELTA_ANGLE_THRESHOLD;
                bool isASmallPanel = (length_width_ratio_panel > ARMOR_LENGTH_WIDTH_RATIO_SMALL * 0.66) 
                                  && (length_width_ratio_panel < ARMOR_LENGTH_WIDTH_RATIO_SMALL * 1.37);
                bool isALargePanel = 0;/*(length_width_ratio_panel > ARMOR_LENGTH_WIDTH_RATIO_LARGE * 0.66) 
                                  && (length_width_ratio_panel < ARMOR_LENGTH_WIDTH_RATIO_LARGE * 1.10);*/
                if (isDeltaAngleValid && (isASmallPanel || isALargePanel) && isTwoEdgeAtSameLevel)
                {
                    Armor temp = edge_to_armor(edges_detected[i], edges_detected[j], distance);
                    armor_detected.push_back(temp);
                }
            }
        }
        return true;
    }
}

Armor Armor_recog::edge_to_armor(cv::RotatedRect &a, cv::RotatedRect &b, double width)
{
    Point centerA = a.center, centerB = b.center;
    Point2f center = (centerA + centerB) / 2.0;
    Armor armor;
    double widthA = min(a.size.width, a.size.height);
    double lengthA = max(a.size.width, a.size.height);
    double widthB = min(b.size.width, b.size.height);
    double lengthB = min(b.size.width, b.size.height);
    float angle = std::atan2(centerA.y - centerB.y, centerA.x - centerB.x);
    float height = max(lengthA, lengthB);
    
    armor.bound = RotatedRect(center, Size2f(width, height), angle * 180 / CV_PI);
    armor.center = center;
    armor.barLen = (lengthA + lengthB) / 2;
    armor.distance = abs(a.center.x - b.center.x);
    armor.now = this->now;
    return armor;
}

void Armor_recog::showImage()
{
    Mat testa;
    Mat testb;
    current_frame.copyTo(testa);
    current_frame.copyTo(testb);

    for (size_t i = 0; i < edges_detected.size(); i++)
    {
        Point2f vertices[4];
        edges_detected[i].points(vertices);
        for (int j = 0; j < 4; j++)
        {
            line(testa, vertices[j], vertices[(j + 1) % 4], Scalar(255, 0, 0), 1);
        }
    }
    Point2d a(300, 200);
    Point2d b(200, 300);
    line(testa, a, b, Scalar(0, 0, 255));
#ifdef DEBUG_MODE
    imshow("bar", testa);
#endif
    for (size_t i = 0; i < armor_detected.size(); i++)
    {
        Point2f vertices[4];
        armor_detected[i].bound.points(vertices);
        for (int j = 0; j < 4; j++)
        {
            line(testb, vertices[j], vertices[(j + 1) % 4], Scalar(255, 0, 0), 3);
        }
    }
#ifdef DEBUG_MODE
    imshow("armor", testb);
#endif
}

void Armor_recog::detect_armor(Mat &src)
{
    clean();
    /*ROI*/
    Mat temp;
    // std::chrono::time_point<std::chrono::steady_clock> now = std::chrono::steady_clock::now();
    src.copyTo(temp);
    if (!to_aim.isEmpty())
    {
        //previous frame detects armor
        //apply ROI: region of interest method
        Rect rect = to_aim.bound.boundingRect();
        int w = int(rect.width * 2);
        int h = int(rect.height * 2);
        rect = Rect(rect.x, rect.y, w, h);

        try
        {
            temp = src(rect);
        }
        catch (Exception e)
        {
            temp = src;
        }
    }
    if (!read_frame(temp))
    {
        cerr << "No Frame read in. From Armor_recog::detect_armor" << endl;
    }
    pre_process();
    edge_detect();
    armor_detect();
    showImage();
    // waitKey(1);
}

void Armor_recog::clean()
{
    armor_detected.clear();
    edges_detected.clear();
}

bool Armor_recog::detect_rune(Mat &src)
{
    clean();
    int counter = 0;
    vector<vector<Point>> contours;
    Point flood_seed = Point(5, 50);
    read_frame(src);
    pre_processing_rune();


    floodFill(processed_frame, flood_seed, Scalar(255), 0, FLOODFILL_FIXED_RANGE);
    threshold(processed_frame, processed_frame, 80, 255, THRESH_BINARY_INV);
    findContours(processed_frame, contours, RETR_LIST, CHAIN_APPROX_NONE);

    for (size_t i = 0; i < contours.size(); i++)
    {
        vector<Point> points;
        RotatedRect rrect;
        cv::Point2f vertices[4];
        float aim;
        double area = contourArea(contours[i]);

        if (area < 50 || 1e4 < area) continue;

        cv::drawContours(current_frame, contours, static_cast<int>(i), Scalar(0), 2);
        points = contours[i];
        rrect = fitEllipse(points);
        rrect.points(vertices);
        aim = rrect.size.height / rrect.size.width;

        if (aim > 1.7 && aim < 2.6)
        {
            float middle = 1000000.f;
            for (int j = 0; j < 4; j++)
            {
                cv::line(current_frame, vertices[j], vertices[(j + 1) % 4], cv::Scalar(0, 255, 0), 4);
            }

            for (size_t j = 1; j < contours.size(); j++)
            {
                // if(j == i) continue;
                vector<Point> points_next;
                RotatedRect rrectA;
                float height_width_ratio_next;
                double area_next = contourArea(contours[j]);

                if (area_next < 50 || 1e4 < area_next) continue;
                points_next = contours[j];
                rrectA = fitEllipse(points_next);
                height_width_ratio_next = rrectA.size.height / rrectA.size.width;

                if (height_width_ratio_next > 3.0)
                {
                    float distance = sqrt((rrect.center.x - rrectA.center.x) * (rrect.center.x - rrectA.center.x) +
                                          (rrect.center.y - rrectA.center.y) * (rrect.center.y - rrectA.center.y));
                    if (middle > distance) middle = distance;
                }
            }
            if (middle > RUNE_BLOCK_DISTANCE_THRES)
            { 
                // cv::circle(current_frame,Point(rrect.center.x,rrect.center.y),15,cv::Scalar(0,0,255),4);
                Mat prediction = KF.predict();
                Point predict_pt = Point((int)prediction.at<float>(0), (int)prediction.at<float>(1));
                Armor armor;

                measurement.at<float>(0) = (float)rrect.center.x;
                measurement.at<float>(1) = (float)rrect.center.y;
                KF.correct(measurement);
                circle(current_frame, predict_pt, 3, Scalar(34, 255, 255), -1);

                KF.predict(); //迷惑行为
                KF.predict();
                KF.predict();
                KF.predict();
                rrect.center.x = (int)prediction.at<float>(0);
                rrect.center.y = (int)prediction.at<float>(1);
                armor = Armor(rrect, max(rrect.size.width, rrect.size.height));
                armor.center = rrect.center;
                armor_detected.push_back(armor);
                counter++;
            }
        }
    }
    if (counter > 1 || counter == 0)
    {
        for (int i = 0; i < counter; i++)
        {
            armor_detected.pop_back();
        }
        return false;
    }
    else
    {
        return true;
    }
}

void Armor_recog::pre_processing_rune()
{
    resize(current_frame, current_frame, Size(current_frame.cols * 0.5, current_frame.rows * 0.5));
    current_frame.copyTo(processed_frame);
    pre_process();
    dilate(processed_frame, processed_frame, Mat());
}

void Armor_recog::KF_initialize()
{
    int stateNum = 4;
    int measureNum = 2;
    this->KF = KalmanFilter(stateNum, measureNum, 0);
    measurement = Mat::zeros(measureNum, 1, CV_32F);
    this->KF.transitionMatrix = (Mat_<float>(stateNum, stateNum) << 
                           1, 0, 1, 0, //A 状态转移矩阵
                           0, 1, 0, 1,
                           0, 0, 1, 0,
                           0, 0, 0, 1);
    setIdentity(KF.measurementMatrix);                      //H=[1,0,0,0;0,1,0,0] 测量矩阵
    setIdentity(KF.processNoiseCov, Scalar::all(1e-5));     //Q高斯白噪声，单位阵
    setIdentity(KF.measurementNoiseCov, Scalar::all(1e-1)); //R高斯白噪声，单位阵
    setIdentity(KF.errorCovPost, Scalar::all(1));           //P后验误差估计协方差矩阵，初始化为单位阵
    randn(KF.statePost, Scalar::all(0), Scalar::all(0.1));  //初始化状态为随机值
}

bool Armor_recog::setAim()
{
    if (!armor_detected.size())
    {
        this->to_aim = Armor();
        return false;
    }
    sort(armor_detected.begin(), armor_detected.end(), [](Armor &a1, Armor &a2) {
        return a1.bound.size.area() > a2.bound.size.area();
    });
    this->to_aim = armor_detected[0];
    return true;
}

float Armor_recog::con_disjunction_ratio(const Armor &former, const Armor &latter) {
    /*compute the disjunction region over the conjunction region of the two Armor*/
    Point2f up_left_f, up_left_l, down_right_f, down_right_l;
    Point2f up_left_dis, up_left_con, down_right_dis, down_right_con, con;
    float ratio;
    float area_con, area_dis;

    up_left_f = Point(former.bound.center.x - former.bound.size.width/2, 
                      former.bound.center.y - former.bound.size.height/2);
    up_left_l = Point(latter.bound.center.x - latter.bound.size.width/2, 
                      latter.bound.center.y - latter.bound.size.height/2);
    down_right_f = Point(former.bound.center.x + former.bound.size.width/2, 
                      former.bound.center.y + former.bound.size.height/2);
    down_right_l = Point(latter.bound.center.x + latter.bound.size.width/2, 
                      latter.bound.center.y + latter.bound.size.height/2);
    up_left_con = Point(max(up_left_f.x, up_left_l.x), max(up_left_f.y, up_left_l.y));
    down_right_con = Point(min(down_right_f.x, down_right_l.x), min(down_right_f.y, down_right_l.y));
    con = (down_right_con - up_left_con);
    area_con = con.x * con.y;
    area_dis = former.bound.size.width * former.bound.size.height 
              + latter.bound.size.width * latter.bound.size.height - area_con;
    return area_con > 0? area_con/area_dis: 0;
}